#include <sys/ptrace.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <syscall.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>

struct user_pt_regs {
    __u64 regs[31];  // X0~X30
    __u64 sp;        // X31 (SP)
    __u64 pc;        // PC (程序计数器)
    __u64 pstate;    // PSTATE (处理器状态)
};

#define __NR_mmap 222     // ARM64 mmap 系统调用号
#define __NR_munmap 215   // ARM64 munmap 系统调用号

int main() {
    pid_t child = fork();
    if (child == 0) {
        // 子进程：被跟踪
        ptrace(PTRACE_TRACEME, 0, NULL, NULL);
        execve("/bin/sleep", (char*[]){"sleep", "10", NULL}, NULL);
        exit(0);
    }

    int status;
    waitpid(child, &status, 0);  // 等待子进程暂停

    struct user_pt_regs regs;
    struct iovec iov = {
        .iov_base = &regs,
        .iov_len = sizeof(regs)
    };

    // 获取当前寄存器状态
    if (ptrace(PTRACE_GETREGSET, child, (void*)NT_PRSTATUS, &iov) < 0) {
        perror("PTRACE_GETREGSET failed");
        return 1;
    }

    // 设置 mmap 参数
    regs.regs[8] = __NR_mmap;     // x8 = 系统调用号
    regs.regs[0] = 0;             // x0 = addr = NULL
    regs.regs[1] = 4096;          // x1 = length = 4096
    regs.regs[2] = PROT_READ | PROT_WRITE;  // x2 = prot
    regs.regs[3] = MAP_ANONYMOUS | MAP_PRIVATE;  // x3 = flags
    regs.regs[4] = -1;            // x4 = fd
    regs.regs[5] = 0;             // x5 = offset

    // 写入寄存器
    if (ptrace(PTRACE_SETREGSET, child, (void*)NT_PRSTATUS, &iov) < 0) {
        perror("PTRACE_SETREGSET failed");
        return 1;
    }

    // 触发系统调用
    ptrace(PTRACE_SYSCALL, child, NULL, NULL);
    waitpid(child, &status, 0);

    // 获取 mmap 返回值
    if (ptrace(PTRACE_GETREGSET, child, (void*)NT_PRSTATUS, &iov) < 0) {
        perror("PTRACE_GETREGSET failed");
        return 1;
    }
    long mmap_addr = regs.regs[0];
    printf("mmap returned address: %p\n", (void*)mmap_addr);

    // 设置 munmap 参数
    regs.regs[8] = __NR_munmap;
    regs.regs[0] = mmap_addr;
    regs.regs[1] = 4096;

    // 写入寄存器
    if (ptrace(PTRACE_SETREGSET, child, (void*)NT_PRSTATUS, &iov) < 0) {
        perror("PTRACE_SETREGSET failed");
        return 1;
    }

    // 触发系统调用
    ptrace(PTRACE_SYSCALL, child, NULL, NULL);
    waitpid(child, &status, 0);

    // 获取 munmap 返回值
    if (ptrace(PTRACE_GETREGSET, child, (void*)NT_PRSTATUS, &iov) < 0) {
        perror("PTRACE_GETREGSET failed");
        return 1;
    }
    long ret = regs.regs[0];
    printf("munmap returned: %ld\n", ret);

    ptrace(PTRACE_DETACH, child, NULL, NULL);
    return 0;
}